import os
import base64
from PIL import Image
from io import BytesIO
from pathlib import Path
import streamlit.components.v1 as components

_RELEASE = True

if not _RELEASE:
    _component_func = components.declare_component(
        "streamlit_image_gallery",
        url="http://localhost:3001",
    )
else:
    parent_dir = os.path.dirname(os.path.abspath(__file__))
    build_dir = os.path.join(parent_dir, "frontend/build")
    _component_func = components.declare_component(
        "streamlit_image_gallery", path=build_dir
    )


def _is_local_file(path):
    """检查是否为本地文件路径"""
    if not path:
        return False

    # 检查是否为URL
    if path.startswith(("http://", "https://", "//")):
        return False

    # 检查文件是否存在
    return os.path.isfile(path)


def _encode_image_to_base64(image_path: str):
    """将本地图片编码为base64"""
    try:
        with open(image_path, "rb") as image_file:
            encoded_string = base64.b64encode(image_file.read()).decode()

        # 根据文件扩展名确定MIME类型
        ext = Path(image_path).suffix.lower()
        mime_types = {
            ".jpg": "image/jpeg",
            ".jpeg": "image/jpeg",
            ".png": "image/png",
            ".gif": "image/gif",
            ".bmp": "image/bmp",
            ".webp": "image/webp",
        }
        mime_type = mime_types.get(ext, "image/jpeg")

        return f"data:{mime_type};base64,{encoded_string}"
    except Exception as e:
        raise ValueError(f"无法读取图片文件 {image_path}: {str(e)}")


def streamlit_image_gallery(
    images,
    max_width=400,
    gap=4,
    max_cols=2,
    max_rows=2,
    key=None,
    height=150,
    border_color=None,
    clicked_border_color=None,
):
    """Create a new instance of "streamlit_image_gallery".

    The Gallery shows a grid of images. Images can be opened in a separate tab.
    Supports both local files and remote URLs.

    Parameters
    ----------
    images: list of dict
        Image dictionaries: `{ 'src': 'image_url_or_path', 'title': 'Example Title' }`
        For local images, provide the file path (e.g., 'images/photo.jpg')
        For remote images, provide the URL (e.g., 'https://example.com/photo.jpg')
        For PIL.Image, provide the image object (e.g., Image.open('image.png'))
        For base64, best for base64 string (e.g., '...')
    max_width: int or str
        Maximum width of the gallery.
    gap: int
        Gap between the images.
    max_cols: int
        maximum number of columns.
    max_rows: int
        maximum number of rows.
    height: int
        height of the gallery.
    border_color: str or None
        Color of the image borders (e.g., '#ff0000', 'red', 'rgba(255,0,0,0.5)')
    clicked_border_color: str or None
        Color of the clicked image border (e.g., '#ff0000', 'red', 'rgba(255,0,0,0.5)')
    key: str or None
        An optional key that uniquely identifies this component. If this is
        None, and the component's arguments are changed, the component will
        be re-mounted in the Streamlit frontend and lose its current state.
    """
    # 处理图片列表，将本地文件转换为base64
    processed_images = []
    for image in images:
        processed_image = image.copy()

        if image["src"].startswith("data:"):
            # base64图片
            processed_image["is_local"] = True
        elif isinstance(image["src"], Image.Image):
            # 将PIL.Image转换为base64
            buffered = BytesIO()
            image["src"].save(buffered, format="PNG")
            img_str = base64.b64encode(buffered.getvalue()).decode()
            processed_image["src"] = f"data:image/png;base64,{img_str}"
            processed_image["is_local"] = True
        elif _is_local_file(image["src"]):
            # 将本地图片转换为base64
            processed_image["src"] = _encode_image_to_base64(image["src"])
            processed_image["is_local"] = True
        else:
            # 网络图片
            processed_image["is_local"] = False

        processed_images.append(processed_image)

    return _component_func(
        images=processed_images,
        max_width=max_width,
        gap=gap,
        max_cols=max_cols,
        max_rows=max_rows,
        key=key,
        height=height,
        border_color=border_color,
        clicked_border_color=clicked_border_color,
    )


# image list
# images = [
#     {
#         "src": "https://images.unsplash.com/photo-1718439111428-f6ef86aae18d",
#         "title": "flower",
#     },
#     {
#         "src": "https://images.unsplash.com/photo-1718554517666-2978ede88574",
#         "title": "birld",
#     },
#     {
#         "src": "image.png",
#         "title": "local image",
#     },
# ]
# # 显示图片画廊
# result = streamlit_image_gallery(
#     images=images,
#     max_cols=3,
#     gap=10,
#     max_width=800,
#     height=200,
#     key="test_gallery",
#     border_color="gray",
# )
